package com.zx.flow.api.util;

import java.util.UUID;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;

public class TraceUtil {

    public static final String TRACE_ID = "traceId";

    public static final String SPAN_ID = "spanId";

    /**
     * 初始化trace【同步线程场景使用】，同一个线程内允许多次调用，不会影响
     */
    public static Trace initTrace() {
        String traceId = MdcUtil.get(TRACE_ID);
        if (StringUtils.isBlank(traceId)) {
            traceId = TraceUtil.generateTraceId();
        }
        setTraceId(traceId);

        String spanId = MdcUtil.get(SPAN_ID);
        if (StringUtils.isBlank(spanId)) {
            spanId = TraceUtil.generateSpanId();
        }
        setSpanId(spanId);

        Trace trace = new Trace();
        trace.setTraceId(traceId);
        trace.setSpanId(spanId);
        return trace;
    }

    /**
     * 初始化trace，强制指定traceId，并重新分配spanId。【不支持重复调用，每次调用的spanId不同】
     */
    public static Trace initTrace(String traceId) {
        TraceUtil.setTraceId(traceId);
        String spanId = TraceUtil.generateSpanId();
        TraceUtil.setSpanId(spanId);

        Trace trace = new Trace();
        trace.setTraceId(traceId);
        trace.setSpanId(spanId);
        return trace;
    }

    /**
     * 获得trace
     */
    public static Trace getTrace() {
        Trace trace = new Trace();
        trace.setTraceId(TraceUtil.getTraceId());
        trace.setSpanId(TraceUtil.getSpanId());
        return trace;
    }

    public static boolean hasTrace() {
        return !StringUtils.isBlank(TraceUtil.getTraceId());
    }

    /**
     * 清空trace
     */
    public static void clearTrace() {
        MdcUtil.clear();
    }


    /**
     * 写入traceId
     *
     * @return 返回当前的traceId
     */
    public static void setTraceId(String traceId) {
        MdcUtil.put(TRACE_ID, traceId);
    }

    /**
     * 写入spanId
     *
     * @return 返回当前的spanId
     */
    public static void setSpanId(String spanId) {
        MdcUtil.put(SPAN_ID, spanId);
    }

    /**
     * 获得traceId
     */
    public static String getTraceId() {
        return MdcUtil.get(TRACE_ID);
    }

    /**
     * 获得spanId
     */
    public static String getSpanId() {
        return MdcUtil.get(SPAN_ID);
    }

    /**
     * TraceId生成
     *
     * @return TraceId
     */
    public static String generateTraceId() {
        return UUID.randomUUID().toString().replaceAll("-", "").substring(0, 8);
    }

    /**
     * spanId生成
     *
     * @return spanId
     */
    public static String generateSpanId() {
        return UUID.randomUUID().toString().replaceAll("-", "").substring(0, 4);
    }


    @Data
    public static class Trace {

        private String traceId;
        private String spanId;
    }
}
